import json
from functools import cache

from test1 import file_name
from tools.DrawModel import Draw


class OperationalData:
    def __init__(self):
        self.num = 0
        self.car = None
        self.file_name = None
        self.data = None
        self.frame_data = {}
        self.result = []
        self.chart = None

    def load_file(self, file_name):
        self.file_name = file_name
        with open(file_name, 'r') as file:
            self.data = json.load(file)
        self.num = len(self.data)

        for i in self.data:
            self.frame_data[i["iterations"]] = i
        names = list(self.data[0]["parameters"].keys())
        self.car = names

    def get_car_name(self):
        return self.car

    def get_track_data(self, name, key, begin_frame: int=None, end_frame: int=None):
        data = []
        for i in self.frame_data.keys():
            if begin_frame is not None and i < begin_frame:
                continue
            if end_frame is not None and i > end_frame:
                continue
            data.append(self.frame_data[i]["track"][name]["current_state"][key])
        return data

    def get_state_data(self, name, key, frame):
        return self.frame_data[frame]["state"][name]["state"][key]

    def get_frame_data(self, frame):
        return self.frame_data[frame]

    def drawRoad(self, name: list, frame: int=1, filename: str=None):
        if self.chart is None:
            Road = Draw(length=60, width=15)
            Road.ax.set_xlim(70, 170)
            Road.initRoad()
        else:
            Road = self.chart
        for i in name:
            Road.drawTrajectory(x=self.frame_data[frame]["state"][i]["state"]["x"], y=self.frame_data[frame]["state"][i]["state"]["y"],
                                h=self.frame_data[frame]["state"][i]["state"]["h"])
            Road.drawCar(begin_x=self.frame_data[frame]["state"][i]["state"]["x"][0], begin_y=self.frame_data[frame]["state"][i]["state"]["y"][0],
                         heading=self.frame_data[frame]["state"][i]["state"]["h"][0], car_name=i)

        data = {
            "name": "trajectory",
            "chart": Road
        }
        self.chart = Road
        if filename is None:
            Road.savefig(data["name"])
        else:
            Road.savefig(filename)

    def drawState(self, name: list, frame: int=1, filename: str=None):
        if self.chart is None:

            Road = Draw(length=30, width=10)
            Road.ax.set_xlim(5, 55)
            Road.initRoad()
        else:
            Road = self.chart
        if len(name) == 0:
            car = self.car
        else:
            car = name
        for i in car:
            Road.drawCar(begin_x=self.frame_data[frame]["track"][i]["current_state"]["x"],
                         begin_y=self.frame_data[frame]["track"][i]["current_state"]["y"],
                         heading=self.frame_data[frame]["track"][i]["current_state"]["h"], car_name=i)
        data = {
            "name": "state",
            "chart": Road
        }
        self.chart = Road

        if filename is None:
            Road.savefig(data["name"])
        else:
            Road.savefig(filename)

t = OperationalData()
t.load_file("/Users/xutao/Desktop/项目/个人项目/GraduationDesignTool/0_400__nores/2025-02-20_09_38_49.json")
frame = 1
filename = f"big2_duilie0_10_{str(frame)}"
# t.drawRoad(name=["Car5"], filename="Car3biandao", frame=frame)
# t.drawRoad(name=["Car4"], filename="Car3biandao", frame=frame)
# t.drawRoad(name=["Car5"], filename=filename, frame=frame)
# t.drawRoad(name=["Car0"], filename=filename, frame=frame)
# t.drawState(["Car3", "Car4", "Car5"],filename=filename, frame=frame)
# t.drawState(["Car0", "Car1", "Car2"],filename=filename, frame=frame)
t.drawState([],filename=filename, frame=frame)



